import sys
#import pipeline
import numpy as np
import markdown , time , json , os

from PyQt5 import uic
from PyQt5.QtWidgets import QApplication , QWidget , QFileDialog , QPushButton , \
    QLabel , QVBoxLayout , QHBoxLayout , QListWidgetItem  , QCheckBox
from PyQt5.QtGui import QFont
from PyQt5.QtCore import pyqtSignal , QThread , Qt
from SubThread import * 
from bubble_message import ImageMessage , MessageType, Notice , BubbleMessage
from image_process import ImageWidget
from MyListModel import * 
from htmlEdit import insert_keyword , format_think_content

# warnings.filterwarnings('ignore',category = DeprecationWarning) #屏蔽警告

class MainWindow_0(QWidget):
    def __init__(self):
        super().__init__()
        
        self.initUI()                       # 初始化ui界面

        self.communicateRecord = np.array([])
        self.answer = None
        self.__slotConnect()                # 链接主线程的槽函数
        self.__subThreadInit()              # 初始化子线程
        # self.__test()                     # 测试函数

        # self.__chatWidgetTest()           # 测试聊天窗口正不正常

    def initUI(self):                       # 初始化ui界面
        self.ui = uic.loadUi("ui\MainWindow.ui",self)
        self.ui_imageWindow = ImageWidget()
        # print(type(self.ui))
        self.ui.chatWidget.setAvatar(       # 设置聊天双方的头像
            send = r'data\head4.jpg',
            receive = r'data\head3.jpg'
        )

        # 时间
        self.time_message = Notice(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()))
        self.ui.chatWidget.add_message_item(self.time_message)

        self.ui.ButtonStopGenerate.setVisible(False)

        self.updateHistoryList()

    def __test(self):
        # btn1 = QPushButton()
        # btn1.setText('btn')

        # layout = QVBoxLayout()
        # self.ui.verticalLayout = QHBoxLayout()
        # print(type(self.ui.chatWidget))
        
        # self.ui.chatWidget.layout0.addWidget(btn1)
        # self.ui.setLayout(layout)
        pass

    def __chatWidgetTest(self):             # 测试用，用来测试聊天窗口正不正常
        # self.ui.chatWidget = ChatWidget()
        TEXT = MessageType.Text
        IMAGE = MessageType.Image
        time_message = Notice('2023-11-18 17:43')
        self.ui.chatWidget.add_message_item(time_message)
        bubble_message = BubbleMessage('你好think你好你好你好你好你好你好你好你好你好你好你好你好你好你好你好你好', \
                                       self.ui.chatWidget.receive_avatar, Type=TEXT, is_send=False)
        self.ui.chatWidget.add_message_item(bubble_message)
        bubble_message = BubbleMessage('你好啊💖', self.ui.chatWidget.send_avatar, Type=TEXT, is_send=True)
        self.ui.chatWidget.add_message_item(bubble_message)

    def __slotConnect(self):                # 链接一般槽函数
        self.ui.pushButtonUTr.clicked.connect(self.SignalUserTransmit)
        self.ui.ButtonStopGenerate.clicked.connect(self.stopGenerate)
        self.ui.ButtonClearChatWidget.clicked.connect(self.clearChatWidget)
        self.ui.ButtonSaveChatHistory.clicked.connect(self.slotSaveChatHistory)
        self.ui.ButtonImageUpload.clicked.connect(self.imageUploadCmd)
        self.ui.ButtonLoadHistory.clicked.connect(self.slotLoadChatHistory)
        self.ui.ButtonLoadPdf.clicked.connect(self.slotLoadPdf)
        self.ui.HistoryListWidget.itemDoubleClicked.connect(self.slotSelectChatHistory)
        self.ui_imageWindow.confirmButton.clicked.connect(self.slotImageSelectConfirm)
        self.ui_imageWindow.reSelectButton.clicked.connect(self.imageUploadCmd)

    def __subThreadInit(self):                          # 子线程初始化函数
        self.getResponse_Thread = getLangChainResponse_Thread()  # 初始化获取回答的子线程
        self.UpLoadImage_Thread = UpLoadImage_Thread()  # 初始化图像识别的子线程
        # self.AudioInput_Thread = AudioInput_Thread()    # 初始化音频输入的子线程

        self.getResponse_Thread.StreamResponseSignal.connect(self.UpdateBotResponse)    # 链接自定义信号和槽函数
        self.UpLoadImage_Thread.imageWidgetCmdSignal.connect(self.subThreadOCRResultProcess)
        self.getResponse_Thread.finishSignal.connect(self.ResponseFinish)
        self.UpLoadImage_Thread.finishSignal.connect(self.ImageUpLoadFinish)

    def keyPressEvent(self, event):                         # 键盘按键触发事件
        key = event.key()
        modifiers = event.modifiers()
        if key == Qt.Key_S and modifiers == Qt.ControlModifier:
            self.slotSaveChatHistory()
            print("Ctrl+S pressed")
        elif (key == Qt.Key_Enter or key == Qt.Key_Return) and modifiers == Qt.ControlModifier:
            if self.ui_imageWindow.isVisible():
                self.slotImageSelectConfirm()
            else:
                self.SignalUserTransmit()                       # 回车键快捷发送
            print("Ctrl + Enter pressed")
        elif event.key() == Qt.Key_Escape:
            self.close()
        return super().keyPressEvent(event)
    
    def updateHistoryList(self):
        # 加载历史记录列表：
        try:
            fileNames = os.listdir('./history')
            for name in fileNames:
                self.ui.HistoryListWidget.addItem(name)
        except Exception as e:
            print(f'加载历史记录列表失败，产生的错误为：{e}')

    def SignalUserTransmit(self):                           # 用户提问后获取回答
        userInText = self.ui.textEditUserIn.toPlainText()   # 提取用户输入的文本
        self.ui.textEditUserIn.clear()                      # 清空输入框

        self.ui.chatWidget.appendPlainText(str(userInText),is_send=True) # 更新记录
        # self.ui.chatWidget.addTextBubble(userInText,is_send = True)
        
        self.communicateRecord = np.append(self.communicateRecord,tuple(('user',userInText)))

        self.ui.chatWidget.appendPlainText('思考中，请稍候。。。')
        # self.ui.chatWidget.addTextBubble('')

        self.ui.ButtonStopGenerate.setVisible(True)
        self.ui.pushButtonUTr.setEnabled(False)

        self.getResponse_Thread.is_online = self.ui.SearchOnline.isChecked()
        self.getResponse_Thread.prompt = userInText
        self.getResponse_Thread.start()         # 启动线程，执行线程类中run函数，获取deepseek的流式返回值

    def UpdateBotResponse(self,chunk):          # 更新响应，如果大模型生成了文本就会触发这个函数执行更新操作
        if(chunk == '</think>'):
            # htmlText = markdown.markdown(self.ui.chatWidget.bubbleList[-1].message.text())
            # self.ui.chatWidget.bubbleList[-1].message.setText(htmlText)
            self.ui.chatWidget.appendPlainText('思考中，请稍候。。。')
            self.ui.chatWidget.bubbleList[-1].message.setFont(QFont('微软雅黑', 12))
        elif(chunk == '<think>'):
            self.ui.chatWidget.bubbleList[-1].message.setFont(QFont('楷体', 12))
            self.ui.chatWidget.insertBubbleText('* 思考过程：')
        else:
            self.ui.chatWidget.insertBubbleText(chunk)  # 往对话气泡里插入新生成的文本

    def imageUploadCmd(self):                           # 用户在本地图片中选择图片上传
        # print('image upload button is pressed!' + str(self.UpLoadImage_Thread.isRunning()))
        self.UpLoadImage_Thread.img_path = self.ui_imageWindow.openImage()[0]
        if(not self.ui_imageWindow.isVisible()):
            self.ui_imageWindow.show()
        if(not self.UpLoadImage_Thread.isRunning()):    # 打开图像识别线程
            self.UpLoadImage_Thread.start()

    def slotSelectChatHistory(self,selItem:QListWidgetItem):# 双击某条历史记录的槽函数
        try:
            with open('./history/' + selItem.text(),'r') as file_obj:
                loadDict = json.load(file_obj)
                self.time_message = Notice(loadDict['time_message'])
                self.ui.chatWidget.flush(loadDict['chat_history'])
                self.ui.chatWidget.add_message_item(self.time_message,index = 0)
                # self.ui.chatWidget.flush(self.getResponse_Thread.conversation_list)
                
        except Exception as e:
            print(f"加载历史记录失败，产生的异常为:{e}")

    def slotLoadChatHistory(self):                      # 加载用户选择聊天的历史记录
        try:

            pass
        except Exception as e:
            print(f"加载历史记录失败，产生的异常为{e}")

    def slotLoadPdf(self):
        try:
            fname , ftype = QFileDialog.getOpenFileName(self, "Open File", "", 
                                             "PDF Files(*.pdf);;All Files (*)")#如果添加一个内容则需要加两个分号
            # self.getResponse_Thread.llm.loadDocument(fname)# 加载文件
            # self.getResponse_Thread.doc_list.append(fname)
            # print(f"文档{fname}加载完成")
            self.docInput_Thread = DocInput_Thread(
                getResponse_Thread = self.getResponse_Thread,
                fname = fname)
            self.ui.chatWidget.appendPlainText('正在加载文档，请稍候。。。')
            self.ui.pushButtonUTr.setEnabled(False)
            self.ui.ButtonLoadPdf.setEnabled(False)
            self.docInput_Thread.finishSignal.connect(lambda text:print(text))
            self.docInput_Thread.finishSignal.connect(lambda:self.ui.pushButtonUTr.setEnabled(True))
            self.docInput_Thread.finishSignal.connect(lambda:self.ui.ButtonLoadPdf.setEnabled(True))
            self.docInput_Thread.finishSignal.connect(lambda:self.ui.chatWidget.setBubbleText(f'文档{fname}已加载完毕'))
            self.docInput_Thread.start()


        except Exception as e:
            print(f"加载文档失败，产生的异常为{e}")

    def subThreadOCRResultProcess(self,result) -> str:  # 处理对图像ocr识别的结果 
        self.ui_imageWindow.OCRText = ''
        for iter in result[-1]:
            self.ui_imageWindow.OCRText = self.ui_imageWindow.OCRText + iter[1][0] + '\n'
            # print(str(iter[1][0]))
            if(self.ui_imageWindow.isVisible):
                self.ui_imageWindow.generateRectOnText(iter[0][0][0],iter[0][0][1],\
                                                        iter[0][1][0] - iter[0][0][0],iter[0][2][1] - iter[0][0][1])
                # self.ui_imageWindow.update()    # 更新，触发paintEvent

        print(self.ui_imageWindow.OCRText)
        return self.ui_imageWindow.OCRText

    def slotImageSelectConfirm(self):           # 确定选择图像以后
        self.ui_imageWindow.hide()

        userInputText = self.ui_imageWindow.userTextInput.toPlainText()

        imageMessageBubble = BubbleMessage(self.ui_imageWindow.fname, \
                                           self.ui.chatWidget.send_avatar, \
                                       Type=MessageType.Image, is_send=True)    # 更新聊天窗口（插入图片）
        self.ui.chatWidget.bubbleList.append(imageMessageBubble)
        self.ui.chatWidget.add_message_item(self.ui.chatWidget.bubbleList[-1])

        if userInputText == '':
            combined_prompt = (
                f"用户的图片提问里的文字: {self.ui_imageWindow.OCRText}\n"
                f"请结合用户的图片提问里的文字，分析用户想问的问题，并对这个问题做出回答。"
                f"请注意，分析用户想问的问题应该能清楚地描述用户的意图，并包括图片提问里的文字中的相关信息。"
                f"在你给出的答案中，不要出现你分析出来的问题内容，只需要给出你对这个问题的解答结果即可。"
            )
        else:
            combined_prompt = (
                f"用户的文字提问: {userInputText}\n"
                f"用户的图片提问里的文字: {self.ui_imageWindow.OCRText}\n"
                f"请结合上述两个信息，分析用户想问的问题，并对这个问题做出回答。"
                f"请注意，分析用户想问的问题应该能清楚地描述用户的意图，并包括图片提问里的文字中的相关信息。"
                f"在你给出的答案中，不要出现你分析出来的问题内容，只需要给出你对这个问题的解答结果即可。"
            )
            
            self.ui.chatWidget.appendPlainText(userInputText,is_send=True) # 更新聊天框（插入文字）：
            # self.ui.chatWidget.addTextBubble(userInText,is_send = True)
        
        self.communicateRecord = np.append(self.communicateRecord,tuple(('user',combined_prompt)))

        self.ui.chatWidget.appendPlainText('思考中，请稍候。。。')
        # self.ui.chatWidget.addTextBubble('')

        self.ui.ButtonStopGenerate.setVisible(True)

        self.getResponse_Thread.prompt = combined_prompt
        self.getResponse_Thread.start()         # 启动线程，执行线程类中run函数，获取deepseek的流式返回值

    def ImageUpLoadFinish(self):                # 图像上传线程退出槽函数
        self.UpLoadImage_Thread.quit()
    
    def ResponseFinish(self,answer):            # 大模型响应结束的槽函数
        self.ui.pushButtonUTr.setEnabled(True)

        htmlText = markdown.markdown(self.ui.chatWidget.bubbleList[-1].message.text())
        htmlText = insert_keyword(htmlText,'<think',' style = "display: block; font-size: 16px; color: gray;"',
                                  is_after = True)
        # htmlText = r'''
        '''
        <!DOCTYPE html>
        <html lang="en">
        <head>
            <style>
                think {
                    display: block;
                    font-size: 20px;
                    color: gray;
                    font-family: '宋体'
                    font-style: italic;
                }
                think * {
                    font-size: inherit;
                    color: inherit;
                    font-family: 'inherit';
                    font-style: inherit;
                }
            </style>
        </head>
        ''' # '<body>' + insert_keyword(htmlText,'<think>','<p>',is_after = True) + '</body>'
        # htmlText = insert_keyword(htmlText,'</think>','</p>')
        
        # htmlText = format_think_content(htmlText)
        print(htmlText)
        self.ui.chatWidget.bubbleList[-1].message.setText(htmlText)
        self.communicateRecord = np.append(self.communicateRecord,tuple(('bot',answer)))
        # print(self.communicateRecord)     # 打印记录，调试用 

        self.ui.ButtonStopGenerate.setVisible(False)
        self.getResponse_Thread.quit()

    def stopGenerate(self):
        if self.getResponse_Thread.isRunning():
            self.getResponse_Thread.requestInterruption()   # 请求中断子线程

    def slotSaveChatHistory(self):                          # 保存聊天历史记录的槽函数
        try:
            # 以下为保存聊天记录到指定文件夹
            # filename = f".\history\{self.time_message.text()}.json"
            myFileDialog = QFileDialog(directory = r'.\history')
            fileDir = myFileDialog.getExistingDirectory()
            print('保存文件路径:' + fileDir)
            print(self.time_message.text())
            filename = fileDir + '/' + self.time_message.text().replace(':','_') + '.json'
            saveDict = {'time_message':self.time_message.text(),'chat_history':self.getResponse_Thread.conversation_list}
            filename.replace(' ','_')
            with open(filename,'w') as file_obj:
                json.dump(saveDict,file_obj)

            # 更新聊天记录文件名表：
            self.updateHistoryList()
            print("save chat history success!")

        except Exception as e:
            print("exception occurred during saving chat history.\t"
                f"the exception is {e}.")

    def clearChatWidget(self):              # 清空聊天区
        if self.ui.chatWidget.bubbleList != []:
            count = len(self.ui.chatWidget.bubbleList)
            for i in range(count):
                self.ui.chatWidget.layout0.removeWidget(self.ui.chatWidget.bubbleList[i])
            self.ui.chatWidget.bubbleList.clear()
            self.communicateRecord = np.array([])
            self.getResponse_Thread.conversation_list.clear()

    def UpdateTextCommunicate(self):        # 更新对话框内容
        pass

if __name__ == '__main__':
    app = QApplication(sys.argv)
    Window_0 = MainWindow_0()
    Window_0.ui.show()
    app.exec_()
